{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp predefined_problems.test_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-defined Problems\n",
    "\n",
    "This is for test purpose. Please don't use them. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import numpy as np\n",
    "from typing import List\n",
    "import inspect\n",
    "\n",
    "from m3tl.preproc_decorator import preprocessing_fn\n",
    "from m3tl.utils import get_or_make_label_encoder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def _gen_fake_text(tokenized=False) -> List[str]:\n",
    "    if not tokenized:\n",
    "        return ['this is a test' for _ in range(10)]\n",
    "    return ['this is a test'.split() for _ in range(10)]\n",
    "\n",
    "\n",
    "def _gen_fake_array() -> np.ndarray:\n",
    "    return [np.random.uniform(size=(10)) for _ in range(10)]\n",
    "\n",
    "\n",
    "def _gen_fake_cate() -> str:\n",
    "    return [0 for _ in range(10)]\n",
    "\n",
    "\n",
    "def _gen_label(label_format='string'):\n",
    "    if label_format == 'string':\n",
    "        return ['a' for _ in range(5)] + ['b' for _ in range(5)]\n",
    "\n",
    "    if label_format == 'float':\n",
    "        return [0.1 for _ in range(10)]\n",
    "\n",
    "    if label_format == 'seq_tag_string':\n",
    "        return ['a b c d'.split() for _ in range(10)]\n",
    "\n",
    "    if label_format == 'array':\n",
    "        return [np.random.uniform(size=(10)) for _ in range(10)]\n",
    "\n",
    "    raise ValueError('Unknown label format: {}'.format(label_format))\n",
    "\n",
    "\n",
    "def generate_fake_data(output_format='list_tuple', label_format='string') -> dict:\n",
    "    if output_format == 'list_tuple':\n",
    "        return _gen_fake_text(tokenized=(label_format == 'seq_tag_string')), _gen_label(label_format=label_format)\n",
    "\n",
    "    if output_format == 'gen_list_tuple':\n",
    "        return zip(\n",
    "            _gen_fake_text(tokenized=(label_format == 'seq_tag_string')), _gen_label(\n",
    "                label_format=label_format)\n",
    "        )\n",
    "\n",
    "    if output_format == 'dict_tuple':\n",
    "        return ({\n",
    "            'inputs_text': _gen_fake_text(),\n",
    "            'inputs_array': _gen_fake_array(),\n",
    "            'inputs_cate': _gen_fake_cate(),\n",
    "            'inputs_cate_modal_type': ['category' for _ in range(10)],\n",
    "            'inputs_cate_modal_info': [1 for _ in range(10)]\n",
    "        },\n",
    "        _gen_label(label_format=label_format))\n",
    "\n",
    "\n",
    "    if output_format == 'gen_dict_tuple':\n",
    "        dict_list = []\n",
    "        for text, arr, cate, label in zip(\n",
    "                _gen_fake_text(),\n",
    "                _gen_fake_array(),\n",
    "                _gen_fake_cate(),\n",
    "                _gen_label(label_format=label_format)):\n",
    "            dict_list.append(({\n",
    "                'inputs_text': text,\n",
    "                'inputs_array': arr,\n",
    "                'inputs_cate': cate,\n",
    "                'inputs_cate_modal_type': 'category',\n",
    "                'inputs_cate_modal_info': 1,\n",
    "            }, label))\n",
    "        return iter(dict_list)\n",
    "\n",
    "    dict_list = []\n",
    "    for record_id, (text, arr, cate, label) in enumerate(\n",
    "            zip(\n",
    "                _gen_fake_text(),\n",
    "                _gen_fake_array(),\n",
    "                _gen_fake_cate(),\n",
    "                _gen_label(label_format=label_format)\n",
    "            )):\n",
    "        dict_list.append({\n",
    "            'inputs_record_id': record_id,\n",
    "            'inputs_text': text,\n",
    "            'inputs_array': arr,\n",
    "            'inputs_cate': cate,\n",
    "            'inputs_cate_modal_type': 'category',\n",
    "            'inputs_cate_modal_info': 1,\n",
    "            'labels': label\n",
    "        })\n",
    "\n",
    "    if output_format == 'dict':\n",
    "        return dict_list\n",
    "\n",
    "    if output_format == 'gen_dict':\n",
    "        return iter(dict_list)\n",
    "\n",
    "    if output_format == 'rdd':\n",
    "        from pyspark import SparkContext\n",
    "        sc: SparkContext = SparkContext.getOrCreate()\n",
    "        return sc.parallelize(dict_list)\n",
    "\n",
    "\n",
    "def create_fake_label_encoder(params, problem: str, mode: str, label_format='string'):\n",
    "    label_list = _gen_label(label_format=label_format)\n",
    "    get_or_make_label_encoder(\n",
    "        params=params, problem=problem, mode=mode, label_list=label_list, overwrite=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def get_fake_contrastive_learning_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def fake_contrastive_learning(params, mode):\n",
    "        \"\"\"Just a test problem to test multiproblem support\n",
    "\n",
    "        Arguments:\n",
    "            params {Params} -- params\n",
    "            mode {mode} -- mode\n",
    "        \"\"\"\n",
    "        return generate_fake_data(output_format='list_tuple', label_format='string')\n",
    "\n",
    "    return fake_contrastive_learning\n",
    "\n",
    "def get_weibo_fake_cls_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_fake_cls(params, mode):\n",
    "        \"\"\"Just a test problem to test multiproblem support\n",
    "\n",
    "        Arguments:\n",
    "            params {Params} -- params\n",
    "            mode {mode} -- mode\n",
    "        \"\"\"\n",
    "        return generate_fake_data(output_format='list_tuple', label_format='string')\n",
    "\n",
    "    return weibo_fake_cls\n",
    "\n",
    "\n",
    "def get_weibo_fake_ner_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_fake_ner(params, mode):\n",
    "        return generate_fake_data(output_format='list_tuple', label_format='seq_tag_string')\n",
    "    return weibo_fake_ner\n",
    "\n",
    "\n",
    "def get_weibo_pretrain_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_pretrain(params, mode):\n",
    "        return _gen_fake_text()\n",
    "    return weibo_pretrain\n",
    "\n",
    "\n",
    "def get_weibo_fake_multi_cls_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_fake_multi_cls(params, mode):\n",
    "        return generate_fake_data(output_format='dict_tuple', label_format='seq_tag_string')\n",
    "    return weibo_fake_multi_cls\n",
    "\n",
    "\n",
    "def get_weibo_masklm(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_masklm(params, mode):\n",
    "        # Gen has to manual call get or make label encoder!\n",
    "        create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='string')\n",
    "        return generate_fake_data(output_format='gen_list_tuple', label_format='string')\n",
    "    return weibo_masklm\n",
    "\n",
    "\n",
    "def get_weibo_premask_mlm(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_premask_mlm(params, mode):\n",
    "        create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='string')\n",
    "        return generate_fake_data(output_format='gen_list_tuple', label_format='string')\n",
    "    return weibo_premask_mlm\n",
    "\n",
    "\n",
    "def get_weibo_fake_regression(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_fake_regression(params, mode):\n",
    "        return generate_fake_data(output_format='gen_dict', label_format='float')\n",
    "    return weibo_fake_regression\n",
    "\n",
    "\n",
    "def get_weibo_fake_vector_fit(file_path):\n",
    "    @preprocessing_fn\n",
    "    def weibo_fake_vector_fit(params, mode):\n",
    "        create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='array')\n",
    "        return generate_fake_data(output_format='gen_list_tuple', label_format='array')\n",
    "    return weibo_fake_vector_fit\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "Adding new problem weibo_cws, problem type: seq_tag\n",
      "Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "Adding new problem weibo_fake_cls, problem type: cls\n",
      "Adding new problem weibo_masklm, problem type: masklm\n",
      "Adding new problem weibo_fake_regression, problem type: regression\n",
      "Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "Adding new problem weibo_premask_mlm, problem type: premask_mlm\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "from m3tl.special_tokens import TRAIN\n",
    "tb = TestBase()\n",
    "fn = get_weibo_fake_multi_cls_fn('')\n",
    "res = fn(params=tb.params, mode=TRAIN)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'text': ['this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test',\n",
       "   'this is a test'],\n",
       "  'array': [array([0.49095954, 0.20541639, 0.49262084, 0.87451381, 0.02458258,\n",
       "          0.65257202, 0.98258836, 0.58114124, 0.58505677, 0.15806585]),\n",
       "   array([0.44947885, 0.03171102, 0.62730992, 0.7299056 , 0.04007493,\n",
       "          0.65133139, 0.60520342, 0.15861278, 0.29668022, 0.33993872]),\n",
       "   array([0.2731397 , 0.01539462, 0.53834935, 0.32880503, 0.53450392,\n",
       "          0.76220641, 0.44569313, 0.78168369, 0.54579642, 0.71716377]),\n",
       "   array([0.12633738, 0.75136335, 0.32964124, 0.64901069, 0.09048094,\n",
       "          0.96317688, 0.47199601, 0.45125637, 0.53099414, 0.51856527]),\n",
       "   array([0.75426461, 0.76001104, 0.22953306, 0.9394916 , 0.37764847,\n",
       "          0.50555047, 0.93132036, 0.47237353, 0.48608123, 0.0177382 ]),\n",
       "   array([0.4976838 , 0.06866748, 0.37725227, 0.12630836, 0.23961906,\n",
       "          0.21206261, 0.60618255, 0.49934502, 0.28125586, 0.98571187]),\n",
       "   array([0.48514133, 0.24224251, 0.88660613, 0.55056173, 0.4166745 ,\n",
       "          0.09398377, 0.46331419, 0.35790425, 0.53340011, 0.50248541]),\n",
       "   array([0.37552715, 0.22720832, 0.5789555 , 0.60622229, 0.29577143,\n",
       "          0.52823125, 0.74576145, 0.22046372, 0.67712443, 0.06888701]),\n",
       "   array([0.07969023, 0.81459337, 0.011728  , 0.23679496, 0.18646781,\n",
       "          0.22561861, 0.89018834, 0.234746  , 0.36101215, 0.19218532]),\n",
       "   array([0.19514571, 0.74129395, 0.10188232, 0.13237138, 0.5208849 ,\n",
       "          0.24629547, 0.86091529, 0.09563499, 0.59220658, 0.58684688])],\n",
       "  'cate': ['fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate',\n",
       "   'fake_cate']},\n",
       " [['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd'],\n",
       "  ['a', 'b', 'c', 'd']])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generate_fake_data(output_format='dict_tuple', label_format='seq_tag_string')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pyspark Problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "@preprocessing_fn\n",
    "def pyspark_fake_cls(params, mode):\n",
    "    \"\"\"Just a test problem to test multiproblem support\n",
    "\n",
    "    Arguments:\n",
    "        params {Params} -- params\n",
    "        mode {mode} -- mode\n",
    "    \"\"\"\n",
    "    create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='string')\n",
    "    return generate_fake_data(output_format='rdd', label_format='string')\n",
    "\n",
    "\n",
    "@preprocessing_fn\n",
    "def pyspark_fake_seq_tag(params, mode):\n",
    "    create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='seq_tag_string')\n",
    "    return generate_fake_data(output_format='rdd', label_format='seq_tag_string')\n",
    "\n",
    "\n",
    "@preprocessing_fn\n",
    "def pyspark_fake_multi_cls(params, mode):\n",
    "    create_fake_label_encoder(params=params, problem=inspect.currentframe(\n",
    "        ).f_code.co_name, mode=mode, label_format='seq_tag_string')\n",
    "    return generate_fake_data(output_format='rdd', label_format='seq_tag_string')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
